#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# Copyright © 2017 Endless Mobile, Inc.
#
# This library is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation; either
# version 2 of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this library; if not, write to the
# Free Software Foundation, Inc., 59 Temple Place - Suite 330,
# Boston, MA 02111-1307, USA.

import argparse
import os
import subprocess
import sys

import gi
gi.require_version('Flatpak', '1.0')
gi.require_version('GLib', '2.0')
gi.require_version('OSTree', '1.0')
gi.require_version('EosUpdaterUtil', '0')
from gi.repository import Flatpak, GLib, OSTree  # noqa
from gi.repository import EosUpdaterUtil  # noqa


class VolumePreparer:
    """
    Class implementing the eos-updater-prepare-volume command line tool.

    This provides a way to copy OSTree refs onto a new repository on a USB
    stick, along with their dependencies (runtimes, extensions, etc.), in a
    layout which will be detected by OstreeRepoFinderMount.

    The code in this class is currently tightly tied to the command line tool.
    """
    EXIT_FAILED = 1
    EXIT_INVALID_ARGUMENTS = 2
    EXIT_RUN_AS_ROOT = 3

    def __init__(self, volume_path, flatpak_refs, quiet=False, debug=False,
                 preserve_permissions=False):
        self.volume_path = volume_path
        self.flatpak_refs = flatpak_refs
        self.quiet = quiet
        self.debug = debug
        self.preserve_permissions = preserve_permissions

        self.sysroot = OSTree.Sysroot.new_default()

    def __run(self, cmd, fatal_error=True):
        """Run cmd locally."""
        print('# {}'.format(cmd))

        try:
            subprocess.check_call(cmd)
        except subprocess.CalledProcessError:
            if fatal_error:
                raise
            else:
                sys.stderr.write('Warning: command failed: {}\n'.format(cmd))

    def __fail(self, exit_status, message):
        """Print an error to stderr and exit with the given error status."""
        assert(exit_status > 0)

        if not self.quiet:
            sys.stderr.write('%s: %s\n' % (sys.argv[0], message))

        sys.exit(exit_status)

    def _validate_flatpak_ref(self, flatpak_ref):
        """Validate a flatpak ref"""
        try:
            Flatpak.Ref.parse(flatpak_ref)
            return True
        except GLib.Error:
            # This could be IOError.NOT_FOUND or IOError.FAILED. Since the
            # latter is generic, don’t try and match a specific error code in
            # case libflatpak changes it in future.
            return False

    def _get_collection_refs_for_commit(self, rev):
        """Read the commit metadata and return (collection_id, ref, checksum)
        tuples, or None."""
        _, repo = self.sysroot.get_repo()
        _, commit, _ = repo.load_commit(rev)
        metadata_dict = GLib.VariantDict.new(commit.get_child_value(0))

        # Use the collection-refs from the commit metadata; the ones
        # configured on the remote might be different
        collection_binding = metadata_dict.lookup_value(
            OSTree.COMMIT_META_KEY_COLLECTION_BINDING, GLib.VariantType('s'))
        if collection_binding is None:
            return None
        else:
            collection_id = collection_binding.get_string()

        ref_binding = metadata_dict.lookup_value(
            OSTree.COMMIT_META_KEY_REF_BINDING, GLib.VariantType('as'))
        if ref_binding is None:
            return None
        else:
            refs = ref_binding.get_strv()

        return [(collection_id, ref, rev) for ref in refs]

    def _get_os_collection_ref_checksums(self):
        """Get the set of collection–ref and checksum 3-tuples for the booted
        OS, or None."""
        deployment = self.sysroot.get_booted_deployment()
        if not deployment:
            if 'EOS_UPDATER_TEST_UPDATER_DEPLOYMENT_FALLBACK' not in \
               os.environ:
                return None

            deployments = self.sysroot.get_deployments()
            if not deployments:
                return None
            deployment = deployments[0]

        booted_commit = deployment.get_csum()
        if not booted_commit:
            return None

        # The booted commit may have more than one associated ref; copy all
        # of them
        # For example: [('com.endlessm.Os', 'os/eos/amd64/eos3.5',
        #                '470ab0359...'),
        #               ('com.endlessm.Os', 'os/eos/amd64/eos3a',
        #                '470ab0359...')]
        return self._get_collection_refs_for_commit(booted_commit)

    def _get_autoinstall_flatpaks(self):
        """
        Read all the autoinstall lists (see eos-updater-flatpak-installer(8))
        and return a set of the flatpaks listed as needing an install or
        upgrade. Load the set irrespective of the state of the autoinstall
        counter on this system, since whichever system the USB drive is used on
        might have a different counter value.
        """
        autoinstalls = set()

        applied_actions = \
            EosUpdaterUtil.flattened_flatpak_ref_actions_from_paths(None)

        for action in applied_actions:
            if action.type != \
               EosUpdaterUtil.FlatpakRemoteRefActionType.UNINSTALL:
                autoinstalls.add(action.ref.ref.format_ref())
            else:
                autoinstalls.discard(action.ref.ref.format_ref())

        return autoinstalls

    def prepare_volume(self):
        # We need to be root in order to read all the files in the OSTree repo
        # (unless we’re running the unit tests). */
        if os.geteuid() != 0 and \
           'EOS_UPDATER_TEST_UPDATER_DEPLOYMENT_FALLBACK' not in os.environ:
            return self.__fail(self.EXIT_RUN_AS_ROOT, 'Must be run as root')

        # Set up.
        try:
            self.sysroot.load()
        except GLib.Error:
            return self.__fail(self.EXIT_FAILED,
                               'OSTree sysroot could not be loaded; '
                               'are you on an OSTree system?')

        # Try to validate the flatpak refs now rather than failing when
        # `flatpak create-usb` is called
        invalid_refs = [ref for ref in self.flatpak_refs
                        if not self._validate_flatpak_ref(ref)]
        if invalid_refs:
            return self.__fail(self.EXIT_INVALID_ARGUMENTS,
                               'Invalid flatpak refs: %s' %
                               ', '.join(invalid_refs))

        # Add the flatpaks that will be installed by
        # eos-updater-flatpak-installer.
        try:
            autoinstall_flatpaks = self._get_autoinstall_flatpaks()
        except GLib.Error:
            return self.__fail(self.EXIT_FAILED,
                               'Failed to list autoinstall flatpaks to add to '
                               'the USB drive')

        os_collection_refs = self._get_os_collection_ref_checksums()
        if not os_collection_refs:
            return self.__fail(self.EXIT_FAILED,
                               'OSTree deployment ref could not be found; '
                               'are you on an OSTree system?')

        # Pass the heavy lifting off to `ostree create-usb` and
        # `flatpak create-usb`
        _, repo = self.sysroot.get_repo()
        repo_path = os.path.realpath(repo.get_path().get_path())
        for os_collection_ref in os_collection_refs:
            self.__run(['ostree', 'create-usb', '--repo', repo_path,
                        '--commit={}'.format(os_collection_ref[2]),
                        self.volume_path, os_collection_ref[0],
                        os_collection_ref[1]])
        # For specified flatpaks, an error is fatal
        for flatpak_ref in self.flatpak_refs:
            self.__run(['flatpak', '--system', 'create-usb',
                        self.volume_path, flatpak_ref],
                       fatal_error=True)
        # For autoinstall flatpaks, an error is a warning. Some may have been
        # manually uninstalled, and the drive can still be used to update some
        # OS versions.
        for flatpak_ref in autoinstall_flatpaks:
            self.__run(['flatpak', '--system', 'create-usb',
                        self.volume_path, flatpak_ref],
                       fatal_error=False)
        # Ensure that this USB drive can later be read by any user and the
        # .ostree directory can be written to by any user
        if not self.preserve_permissions:
            self.__run(['chmod', '755',
                        GLib.build_filenamev([self.volume_path])])
        self.__run(['chmod', '-R', '777',
                    GLib.build_filenamev([self.volume_path, '.ostree'])])


def main():
    """Main entry point to eos-updater-prepare-volume. Handles arguments."""
    parser = argparse.ArgumentParser(
        description='Prepare a USB drive with a copy of the local OSTree '
                    'repository and the specified flatpaks, so it can be '
                    'used to update other machines offline. The repository '
                    'copy will be put in the .ostree/repo directory on the '
                    'USB drive; other files will not be affected.')
    parser.add_argument('--quiet', action='store_const', const=True,
                        help='do not print anything; check exit status '
                             'for success')
    parser.add_argument('--preserve-permissions', dest='preserve_permissions',
                        action='store_const', const=True,
                        help='don’t make the drive world readable (by '
                             'other user IDs)')
    parser.add_argument('volume_path', metavar='VOLUME-PATH', type=str,
                        help='path to the USB drive to prepare')
    parser.add_argument('flatpak_refs', metavar='REF', nargs='*',
                        help='refs of flatpaks to put on the USB drive '
                             '(e.g. app/com.endlessm.wiki_art.en/x86_64/eos3)')

    args = parser.parse_args()

    VolumePreparer(**vars(args)).prepare_volume()


if __name__ == '__main__':
    main()
